package middleware

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"secure-comm/pkg/crypto"

	"github.com/gin-gonic/gin"
)

// SecureCommunication 加密通信中间件
func SecureCommunication() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 检查是否为加密请求
		if c.GetHeader("X-Encrypted") == "true" {
			// 读取原始请求体
			bodyBytes, err := io.ReadAll(c.Request.Body)
			if err != nil {
				c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "Failed to read request body"})
				return
			}

			// 解析加密请求
			var request struct {
				Encrypted string `json:"encrypted"`
			}
			
			if err := json.Unmarshal(bodyBytes, &request); err != nil {
				c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "Invalid request format"})
				return
			}

			// 解密数据
			var decryptedData map[string]interface{}
			if err := crypto.Decrypt(request.Encrypted, &decryptedData); err != nil {
				c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "Decryption failed: " + err.Error()})
				return
			}

			fmt.Printf("[DEBUG] Decrypted data: %+v\n", decryptedData)

			// 将解密后的数据重新编码为JSON并替换请求体
			decryptedJSON, err := json.Marshal(decryptedData)
			if err != nil {
				c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": "Failed to marshal decrypted data"})
				return
			}

			fmt.Printf("[DEBUG] Decrypted JSON: %s\n", string(decryptedJSON))

			// 替换请求体
			c.Request.Body = io.NopCloser(bytes.NewBuffer(decryptedJSON))
			c.Request.ContentLength = int64(len(decryptedJSON))

			// 标记为已解密
			c.Set("isDecrypted", true)
		}

		// 创建响应拦截器
		responseWriter := &responseWriterWrapper{
			ResponseWriter: c.Writer,
			context:        c,
		}
		c.Writer = responseWriter

		c.Next()
		
		// 处理响应
		responseWriter.processResponse()
	}
}

// responseWriterWrapper 包装响应写入器以拦截响应
type responseWriterWrapper struct {
	gin.ResponseWriter
	context    *gin.Context
	body       []byte
	statusCode int
}

func (w *responseWriterWrapper) Write(data []byte) (int, error) {
	// 检查是否需要加密响应
	if w.context.GetHeader("X-Encrypted") == "true" || w.context.GetBool("isDecrypted") {
		// 拦截响应数据
		w.body = append(w.body, data...)
		return len(data), nil
	}
	
	// 非加密响应直接写入
	return w.ResponseWriter.Write(data)
}

func (w *responseWriterWrapper) WriteString(s string) (int, error) {
	return w.Write([]byte(s))
}

func (w *responseWriterWrapper) WriteHeader(statusCode int) {
	w.statusCode = statusCode
	
	// 如果不需要加密，直接写入状态码
	if w.context.GetHeader("X-Encrypted") != "true" && !w.context.GetBool("isDecrypted") {
		w.ResponseWriter.WriteHeader(statusCode)
	}
}

// processResponse 处理响应数据
func (w *responseWriterWrapper) processResponse() {
	// 检查是否需要加密响应
	if (w.context.GetHeader("X-Encrypted") == "true" || w.context.GetBool("isDecrypted")) && len(w.body) > 0 {
		fmt.Printf("[DEBUG] processResponse: encrypting response, body length: %d\n", len(w.body))
		
		// 解析原始响应
		var originalResponse interface{}
		if err := json.Unmarshal(w.body, &originalResponse); err != nil {
			fmt.Printf("[ERROR] processResponse: failed to unmarshal response: %v\n", err)
			w.ResponseWriter.Header().Set("Content-Type", "application/json")
			w.ResponseWriter.WriteHeader(http.StatusInternalServerError)
			w.ResponseWriter.Write([]byte(`{"error": "Failed to parse response"}`))
			return
		}

		// 加密响应
		encrypted, err := crypto.Encrypt(originalResponse)
		if err != nil {
			fmt.Printf("[ERROR] processResponse: encryption failed: %v\n", err)
			w.ResponseWriter.Header().Set("Content-Type", "application/json")
			w.ResponseWriter.WriteHeader(http.StatusInternalServerError)
			w.ResponseWriter.Write([]byte(`{"error": "Encryption failed"}`))
			return
		}

		// 构造加密响应
		encryptedResponse := map[string]string{"encrypted": encrypted}
		encryptedJSON, _ := json.Marshal(encryptedResponse)

		// 设置响应头 - 必须在WriteHeader之前
		w.ResponseWriter.Header().Set("X-Encrypted", "true")
		w.ResponseWriter.Header().Set("Content-Type", "application/json")
		w.ResponseWriter.Header().Set("Content-Length", fmt.Sprintf("%d", len(encryptedJSON)))
		fmt.Printf("[DEBUG] processResponse: setting X-Encrypted header to true\n")
		
		// 设置状态码
		statusCode := w.statusCode
		if statusCode == 0 {
			statusCode = http.StatusOK
		}
		w.ResponseWriter.WriteHeader(statusCode)

		// 写入加密响应
		w.ResponseWriter.Write(encryptedJSON)
		fmt.Printf("[DEBUG] processResponse: encrypted response written\n")
	} else {
		fmt.Printf("[DEBUG] processResponse: no encryption needed, body length: %d\n", len(w.body))
		// 对于非加密响应，确保状态码被写入
		if w.statusCode != 0 {
			w.ResponseWriter.WriteHeader(w.statusCode)
		}
		// 如果有body数据，写入原始响应
		if len(w.body) > 0 {
			w.ResponseWriter.Write(w.body)
		}
	}
}
